from kj_d1.get_reply import get_ds_reply

import os

from concurrent.futures.thread import ThreadPoolExecutor
from datetime import datetime
import pandas as pd
from loguru import logger

import shutil


def main():
    logger.info('=====启动====')
    if not os.path.exists('res'):
        os.makedirs('res')

    if not os.path.exists('error'):
        os.makedirs('error')
    else:
        shutil.rmtree('error')
        os.makedirs('error')

    had_code = os.listdir('res')
    if '.Ds_Store' in had_code:
        had_code.remove('.Ds_Store')

    had_code = [x.replace('.csv', '') for x in had_code]

    with open('aim_code.txt', 'r', encoding='utf-8') as f:
        need_code = f.readlines()
    need_code = [x.replace('\n', '') for x in need_code]

    need_code = [x for x in need_code if x not in had_code]
    print(need_code)
    print(f'有{len(need_code)}个需要解析')
    
    # run_max_count = 0
    # run_count = 0

    # logger.info(f'重跑次数{run_max_count - run_count}')

    # # 添加一个参数控制是否使用多线程
    # use_multi_thread = True  # 可以通过配置文件或命令行参数来控制
    
    # while len(need_code) > 0:
    #     logger.info(f'目标code：{need_code}')
        
    #     if use_multi_thread:
    #         # 多线程模式
    #         logger.info("使用多线程模式处理")
    #         with ThreadPoolExecutor(max_workers=min(os.cpu_count() - 1, 3)) as executor:
    #             future = executor.map(get_ds_reply, need_code)
    #             list(future)  # 确保所有任务完成
    #     else:
    #         # 单线程模式
    #         logger.info("使用单线程模式处理")
    #         for code in need_code:
    #             get_ds_reply(code)

    #     if run_count >= run_max_count:
    #         break

    #     run_count += 1

    #     had_code = os.listdir('res')
    #     if '.Ds_Store' in had_code:
    #         had_code.remove('.Ds_Store')
    #     had_code = [x.replace('.csv', '') for x in had_code]

    #     with open('aim_code.txt', 'r', encoding='utf-8') as f:
    #         need_code = f.readlines()

    #     need_code = [x.replace('\n', '') for x in need_code]

    #     need_code = [x for x in need_code if x not in had_code]

    ps = [i for i in os.listdir('res') if i.endswith('.csv')]

    df = pd.DataFrame(None)
    for i in ps:
        d1 = pd.read_csv(f'res/{i}', dtype=str, encoding='utf-8')
        df = pd.concat([df, d1], axis=0)

    df.loc[:, '股票代码'] = df.loc[:, '股票代码'].apply(lambda x: str(x).zfill(6))
    df.to_excel(f'跨境商机_{datetime.now().strftime("%Y%m%d")}.xlsx', index=False)


if __name__ == '__main__':
    main()
